import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import pygame
import sys
from pygame.locals import *
import pygame_gui  # 新增pygame_gui库导入

# 设置中文字体（解决中文显示问题）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 单摆运动方程（无量纲形式）
def pendulum_eq(t, y, beta, f, Omega):
    theta, omega = y
    dtheta_dt = omega
    domega_dt = -2 * beta * omega - np.sin(theta) + f * np.cos(Omega * t)
    return [dtheta_dt, domega_dt]

# 模拟参数设置
beta = 0.05  # 无量纲阻尼系数 0.25
f = 0.0  # 无量纲驱动力振幅 (尝试1.07, 1.15, 1.35, 1.45, 1.47, 1.50)
Omega = 2 / 3  # 无量纲驱动力频率
theta0 = 0.5  # 初始摆角

# 初始化pygame
pygame.init()
clock = pygame.time.Clock()  # 增加时钟
WIDTH, HEIGHT = 1200, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("受驱阻尼单摆运动模拟")

# 字体
font = pygame.font.Font('simkai.ttf', 20)
title_font = pygame.font.Font('simkai.ttf', 24)

# 创建GUI管理器
manager = pygame_gui.UIManager((WIDTH, HEIGHT))

# 创建参数输入窗口元素
input_window = pygame_gui.windows.UIConfirmationDialog(
    rect=pygame.Rect((WIDTH//2-300, HEIGHT//2-150), (600, 300)),  # 增加宽度至600
    manager=manager,
    window_title="Parameter Settings",
    action_long_desc="Please enter simulation parameters:",
    action_short_name="Confirm"
)

# 修改所有UILabel的宽度，从200增加到250，以便完整显示文字标签
beta_label = pygame_gui.elements.UILabel(
    relative_rect=pygame.Rect((20, 40), (250, 30)),
    text="Damping coefficient β:",
    manager=manager,
    container=input_window
)
beta_entry = pygame_gui.elements.UITextEntryLine(
    relative_rect=pygame.Rect((280, 40), (280, 30)),  # 调整位置和宽度
    manager=manager,
    container=input_window
)

f_label = pygame_gui.elements.UILabel(
    relative_rect=pygame.Rect((20, 80), (250, 30)),
    text="Driving force amplitude f:",
    manager=manager,
    container=input_window
)
f_entry = pygame_gui.elements.UITextEntryLine(
    relative_rect=pygame.Rect((280, 80), (280, 30)),  # 调整位置和宽度
    manager=manager,
    container=input_window
)

Omega_label = pygame_gui.elements.UILabel(
    relative_rect=pygame.Rect((20, 120), (250, 30)),
    text="Driving force frequency Ω:",
    manager=manager,
    container=input_window
)
Omega_entry = pygame_gui.elements.UITextEntryLine(
    relative_rect=pygame.Rect((280, 120), (280, 30)),  # 调整位置和宽度
    manager=manager,
    container=input_window
)

theta0_label = pygame_gui.elements.UILabel(
    relative_rect=pygame.Rect((20, 160), (250, 30)),
    text="Initial angle θ₀:",
    manager=manager,
    container=input_window
)
theta0_entry = pygame_gui.elements.UITextEntryLine(
    relative_rect=pygame.Rect((280, 160), (280, 30)),  # 调整位置和宽度
    manager=manager,
    container=input_window
)

# 在参数输入窗口增加总时间长度输入框
t_end_label = pygame_gui.elements.UILabel(
    relative_rect=pygame.Rect((20, 200), (250, 30)),
    text="Total simulation time t_end:",
    manager=manager,
    container=input_window
)
t_end_entry = pygame_gui.elements.UITextEntryLine(
    relative_rect=pygame.Rect((280, 200), (280, 30)),
    manager=manager,
    container=input_window
)

# 设置初始值
beta_entry.set_text(str(beta))
f_entry.set_text(str(f))
Omega_entry.set_text(str(Omega))
theta0_entry.set_text(str(theta0))  # 初始摆角默认值
t_end_entry.set_text(str(300))  # 默认总时间长度为300

# 颜色定义
BACKGROUND = (20, 20, 40)

# 参数输入状态
input_active = True
params_set = False

# 参数输入循环
while input_active and not params_set:
    time_delta = clock.tick(60)/1000.0
    
    for event in pygame.event.get():
        if event.type == QUIT:
            input_active = False
            running = False
            params_set = False
        
        manager.process_events(event)
        
        if event.type == pygame_gui.UI_CONFIRMATION_DIALOG_CONFIRMED:
            # 获取输入参数
            try:
                beta = float(beta_entry.get_text())
                f = float(f_entry.get_text())
                Omega = float(Omega_entry.get_text())
                theta0 = float(theta0_entry.get_text())  # 获取初始摆角
                t_end = float(t_end_entry.get_text())    # 获取总时间长度
                params_set = True
            except ValueError:
                # 无效输入处理
                error_dialog = pygame_gui.windows.UIMessageWindow(
                    rect=pygame.Rect((WIDTH//2-150, HEIGHT//2-75), (300, 150)),
                    manager=manager,
                    html_message="<font color='#FF0000'>错误：请输入有效的数值参数</font>"
                )
    
    manager.update(time_delta)
    screen.fill(BACKGROUND)
    manager.draw_ui(screen)
    pygame.display.flip()

# 如果用户取消了输入或关闭窗口
if not params_set:
    pygame.quit()
    sys.exit()

# 初始条件 [theta, omega]
y0 = [theta0, 0.0]  # 使用用户输入的初始摆角

# 时间范围 (无量纲时间，相当于ω₀t)
t_start = 0
# 使用用户输入的总时间长度
t_span = (t_start, t_end)
t_points = 10000  # 时间点数
t_eval = np.linspace(t_start, t_end, t_points)

# 求解微分方程
sol = solve_ivp(pendulum_eq, t_span, y0, args=(beta, f, Omega),
                t_eval=t_eval, rtol=1e-8, atol=1e-8)

# 提取结果
theta = sol.y[0]
omega = sol.y[1]
t = sol.t

# 颜色定义
PENDULUM_COLOR = (70, 200, 255)
ROD_COLOR = (180, 180, 220)
PIVOT_COLOR = (220, 100, 100)
TEXT_COLOR = (220, 220, 255)
GRID_COLOR = (60, 60, 80)
PLOT_COLOR = (0, 200, 150)
PLOT_BG = (30, 30, 50)

# 单摆参数
pendulum_length = 200  # 像素
pivot_x = WIDTH // 4
pivot_y = HEIGHT // 2

# 图表参数
chart_width = WIDTH // 2 - 50
chart_height = HEIGHT // 2 - 50
theta_t_chart_pos = (WIDTH // 2 + 20, 20)
phase_chart_pos = (WIDTH // 2 + 20, HEIGHT // 2 + 20)


# 图表坐标转换函数
def map_to_chart(x, y, x_min, x_max, y_min, y_max, pos_x, pos_y, width, height):
    """将数据坐标映射到屏幕坐标"""
    chart_x = int(pos_x + (x - x_min) / (x_max - x_min) * width)
    chart_y = int(pos_y + height - (y - y_min) / (y_max - y_min) * height)
    return chart_x, chart_y


# 计算图表范围
theta_min, theta_max = min(theta), max(theta)
omega_min, omega_max = min(omega), max(omega)

# 动画主循环
clock = pygame.time.Clock()
running = True
current_index = 0
speed = 1  # 动画速度（每帧推进的数据点数）
paused = False
history_length = 5000  # 历史数据保留时长

while running:
    for event in pygame.event.get():
        if event.type == QUIT:
            running = False
        elif event.type == KEYDOWN:
            if event.key == K_SPACE:
                paused = not paused
            elif event.key == K_UP:
                speed = min(speed + 1, 50)
            elif event.key == K_DOWN:
                speed = max(speed - 1, 1)
            elif event.key == K_r:
                current_index = 0  # 重置动画

    if not paused:
        current_index = min(current_index + speed, len(t) - 1)

    # 清屏
    screen.fill(BACKGROUND)

    # 绘制标题和参数信息 - 减小行距
    title = title_font.render("受驱阻尼单摆运动模拟", True, (100, 200, 255))
    screen.blit(title, (20, 20))

    params = [
        f"参数设置: β = {beta}, f = {f}, Ω = {Omega:.3f}, θ_0 = {theta0}",
        f"当前时间: t = {t[current_index]:.2f}",
        # f"摆角: θ = {theta[current_index]:.3f} rad",
        # f"角速度: ω = {omega[current_index]:.3f} rad/s",
        f"动画速度: {speed} (↑/↓调整, 空格暂停, R重置)"
    ]

    for i, text in enumerate(params):
        text_surface = font.render(text, True, TEXT_COLOR)
        screen.blit(text_surface, (20, 60 + i * 30))  # 行距从40减小到30

    # 绘制单摆
    current_theta = theta[current_index]
    bob_x = pivot_x + pendulum_length * np.sin(current_theta)
    bob_y = pivot_y + pendulum_length * np.cos(current_theta)

    # 绘制支点
    pygame.draw.circle(screen, PIVOT_COLOR, (pivot_x, pivot_y), 10)

    # 绘制摆杆
    pygame.draw.line(screen, ROD_COLOR, (pivot_x, pivot_y), (bob_x, bob_y), 3)

    # 绘制摆球
    pygame.draw.circle(screen, PENDULUM_COLOR, (int(bob_x), int(bob_y)), 20)

    # 绘制角度指示
    pygame.draw.arc(screen, (150, 150, 200),
                    (pivot_x - 50, pivot_y - 50, 100, 100),
                    -np.pi / 2, -np.pi / 2 + current_theta, 2)

    # 绘制角度-时间图
    pygame.draw.rect(screen, PLOT_BG, (*theta_t_chart_pos, chart_width, chart_height))
    pygame.draw.rect(screen, GRID_COLOR, (*theta_t_chart_pos, chart_width, chart_height), 2)

    # 绘制坐标轴和网格
    for i in range(1, 5):
        pygame.draw.line(screen, GRID_COLOR,
                         (theta_t_chart_pos[0], theta_t_chart_pos[1] + i * chart_height // 5),
                         (theta_t_chart_pos[0] + chart_width, theta_t_chart_pos[1] + i * chart_height // 5),
                         1)

    # 绘制θ-t曲线
    points = []
    # 修改为使用固定的历史数据长度
    start_index = max(0, current_index - history_length)
    for i in range(start_index, current_index + 1, 10):
        x, y = map_to_chart(t[i], theta[i], t_start, t_end, theta_min, theta_max,
                            theta_t_chart_pos[0], theta_t_chart_pos[1],
                            chart_width, chart_height)
        points.append((x, y))

    if len(points) > 1:
        pygame.draw.lines(screen, PLOT_COLOR, False, points, 2)

    # 添加图表标题
    theta_t_title = font.render("角度-时间图 (θ-t)", True, TEXT_COLOR)
    screen.blit(theta_t_title, (theta_t_chart_pos[0] + 10, theta_t_chart_pos[1] - 30))

    # 绘制相位图
    pygame.draw.rect(screen, PLOT_BG, (*phase_chart_pos, chart_width, chart_height))
    pygame.draw.rect(screen, GRID_COLOR, (*phase_chart_pos, chart_width, chart_height), 2)

    # 绘制坐标轴和网格
    for i in range(1, 5):
        pygame.draw.line(screen, GRID_COLOR,
                         (phase_chart_pos[0], phase_chart_pos[1] + i * chart_height // 5),
                         (phase_chart_pos[0] + chart_width, phase_chart_pos[1] + i * chart_height // 5),
                         1)
        pygame.draw.line(screen, GRID_COLOR,
                         (phase_chart_pos[0] + i * chart_width // 5, phase_chart_pos[1]),
                         (phase_chart_pos[0] + i * chart_width // 5, phase_chart_pos[1] + chart_height),
                         1)

    # 绘制相位轨迹
    points = []
    # 修改为使用固定的历史数据长度
    start_index = max(0, current_index - history_length)
    for i in range(start_index, current_index + 1, 10):
        x, y = map_to_chart(theta[i], omega[i], theta_min, theta_max, omega_min, omega_max,
                            phase_chart_pos[0], phase_chart_pos[1],
                            chart_width, chart_height)
        points.append((x, y))

    if len(points) > 1:
        pygame.draw.lines(screen, PLOT_COLOR, False, points, 2)

    # 绘制当前位置点
    if current_index > 0:
        phase_x, phase_y = map_to_chart(theta[current_index], omega[current_index],
                                        theta_min, theta_max, omega_min, omega_max,
                                        phase_chart_pos[0], phase_chart_pos[1],
                                        chart_width, chart_height)
        pygame.draw.circle(screen, (255, 100, 100), (phase_x, phase_y), 5)

    # 添加图表标题
    phase_title = font.render("相位图 (θ-ω)", True, TEXT_COLOR)
    screen.blit(phase_title, (phase_chart_pos[0] + 10, phase_chart_pos[1] - 30))

    # 添加图例 - 减小行距
    legend = [
        "单摆运动:",
        f"当前时间: {t[current_index]:.1f}",
        f"摆角: {theta[current_index]:.3f} rad",
        f"角速度: {omega[current_index]:.3f} rad/s",
    ]

    for i, text in enumerate(legend):
        text_surface = font.render(text, True, (180, 150, 255))
        screen.blit(text_surface, (20, HEIGHT - 150 + i * 25))  # 行距25

    # 添加进度条 - 位置微调
    progress = current_index / len(t)
    pygame.draw.rect(screen, (60, 60, 80), (20, HEIGHT - 40, WIDTH - 40, 20))
    pygame.draw.rect(screen, (100, 200, 255), (20, HEIGHT - 40, (WIDTH - 40) * progress, 20))
    # progress_text = font.render(f"进度: {progress * 100:.1f}%", True, TEXT_COLOR)
    # screen.blit(progress_text, (WIDTH // 2 - 60, HEIGHT - 40))

    pygame.display.flip()
    clock.tick(60)

pygame.quit()
sys.exit()